import numpy as np
import neural_network
import datetime
import cv2
import scipy.io as sio
import tensorflow as tf


def my_single(path):
    # 加载自己手写的单张数字
    image = cv2.imread(path, cv2.IMREAD_GRAYSCALE)                          # 获取灰度图
    image = np.uint8(np.where(image > 100, 0, 255))                         # 灰度二值化处理
    new = cv2.resize(image, (28, 28))                                       # 压缩图像为(28,28)
    new = (new-np.min(new)) / (np.max(new)-np.min(new))                     # 图像归一化处理(数据预处理缩小数据量度)
    new = new.reshape((new.shape[0]*new.shape[1], 1))                       # 将矩阵转换为列向量
    return image, new


def load_data(path):
    # 加载.mat数据集
    data = sio.loadmat(path)
    y = data.get('y')  # (5000,1)
    y = y.reshape(y.shape[0])  # make it back to column vector
    X = data.get('X')  # (5000,400)
    return X, y


def data_process(X, y, normal=True):
    """
    数据预处理：列向量展开，归一化，随机排列，标签向量化
    :param X: 三维特征数据集
    :param y: 一维标签数据集
    :param normal: 是否进行归一化
    :return: 三维特征数据集， 三维标签数据集
    """
    # 列向量展开
    X = X.reshape((X.shape[0], X.shape[1] * X.shape[2]))
    # 归一化处理
    if normal:
        X = (X - np.min(X)) / (np.max(X) - np.min(X))
    # 随机排列
    whole = np.insert(X, 0, y, axis=1)
    np.random.shuffle(whole)
    y = whole[:, 0]
    X = np.delete(whole, 0, axis=1)
    # 标签向量化
    y_matrix = []
    for i in range(0, 10):
        y_matrix.append((y == i).astype(int))
    y_matrix = np.array(y_matrix).T
    # 分割为列向量
    new_X = np.array(X).reshape((len(X), len(X[0]), 1))
    new_y = np.array(y_matrix).reshape((len(y_matrix), len(y_matrix[0]), 1))
    return new_X, new_y


if __name__ == "__main__":

    # ************************************************训练模型***************************************************************
    starttime = datetime.datetime.now()

    # 读取数据集
    # 利用tensorflow2.0直接读取MNIST数据集，如果本地没有该数据集文件则会自动上网下载并保存在~/.keras/datasets里名为下面这个‘path’文件
    (train_X, train_y), (test_X, test_y) = tf.keras.datasets.mnist.load_data(path='mnist.npz')
    train_X, train_y = data_process(train_X, train_y)       # 数据预处理
    test_X, test_y = data_process(test_X, test_y)
    # 训练模型
    # 设置神经网络结构，选择损失函数，这里选择交叉熵损失函数（比平方损失函数好点）
    c = neural_network.Neural_Network([784, 1000, 1000, 10], neural_network.CrossEntropy_J())
    #c.check_gradient(newx[0], newy[0])     # 梯度检查，运行非常慢，运行一会看看结果没错就可以关闭了
    c.train(train_X, train_y)
    print(f"准确率为:{c.accuracy(test_X, test_y)*100}%")        # 计算测试集准确率

    endtime = datetime.datetime.now()
    print(f"模型训练时间:{endtime-starttime}")
    # 保存模型在txt文件中
    c.save_model('model_para.txt')
    """
    # ***********************************************直接加载模型*************************************************************
    
    # *************************读取数据集**************************
    # 利用tensorflow2.0直接读取MNIST数据集，如果本地没有该数据集文件则会自动上网下载并保存在~/.keras/datasets里名为下面这个‘path’文件
    (train_X, train_y), (test_X, test_y) = tf.keras.datasets.mnist.load_data(path='mnist.npz')
    train_X, train_y = data_process(train_X, train_y)       # 数据预处理
    test_X, test_y = data_process(test_X, test_y)
    # **************************加载模型***************************
    c = neural_network.Neural_Network([784, 1568, 1568, 10], neural_network.CrossEntropy_J())
    c.load_model('model_para.txt')
    print(f"准确率为:{c.accuracy(test_X, test_y)*100}%")
    """
    # *********************************************识别自己的手写数字**********************************************************
    # 识别自己手写的数字
    file0 = 'C:/Spyder file/Spyder Project/images/9.png'
    image, digital = my_single(file0)
    result = c.predict(digital)
    print(f"数字预测结果：{np.argmax(result)}")
    print(c.layers[-1].output)
    cv2.imshow('window', image)
    cv2.waitKey(0)
